//
// Created by zhulingfeng on 2020/2/6.
//

#include "binary_tree_traversal.hpp"
#include <iostream>
#include <queue>
#include <stack>

void PreOrderTraversalRecursive(BinaryTreeNode *root, std::vector<int> &res) {
  if (!root) return;
  res.push_back(root->val_);
  PreOrderTraversalRecursive(root->left_, res);
  PreOrderTraversalRecursive(root->right_, res);
}

void InOrderTraversalRecursive(BinaryTreeNode *root, std::vector<int> &res) {
  if (!root) return;
  PreOrderTraversalRecursive(root->left_, res);
  res.push_back(root->val_);
  PreOrderTraversalRecursive(root->right_, res);
}

void PostOrderTraversalRecursive(BinaryTreeNode *root, std::vector<int> &res) {
  if (!root) return;
  PostOrderTraversalRecursive(root->left_, res);
  PostOrderTraversalRecursive(root->right_, res);
  res.push_back(root->val_);
}

void BinaryTreeTraversalIterative(BinaryTreeNode *root,
                                  const BinaryTreeTraversalOrder order,
                                  std::vector<int> &res) {
  res.clear();
  std::stack<BinaryTreeNode *> s;
  BinaryTreeNode *last_visited = nullptr;
  auto pre = BinaryTreeTraversalOrder::pre;
  auto in = BinaryTreeTraversalOrder::in;
  auto post = BinaryTreeTraversalOrder::post;

  while (root || !s.empty()) {
    while (root) {
      if (order == pre) {
        res.push_back(root->val_);
      }
      s.push(root);
      root = root->left_;
    }
    root = s.top();
    if (order == pre || order == in) {
      s.pop();
      if (order == in) {
        res.push_back(root->val_);
      }
      root = root->right_;
    } else if (order == post) {
      if (root->right_ == nullptr || root->right_ == last_visited) {
        res.push_back(root->val_);
        last_visited = root;
        s.pop();
        root = nullptr;
      } else {
        root = root->right_;
      }
    } else {
      std::cout << "BinaryTreeTraversalOrder not support: "
                << static_cast<int>(order) << std::endl;
    }
  }
}

void LevelOrderTraversalIterative(BinaryTreeNode *root, std::vector<int> &res) {
  res.clear();
  if (root == nullptr) return;
  std::queue<BinaryTreeNode *> q;
  q.push(root);
  while (!q.empty()) {
    BinaryTreeNode *node = q.front();
    q.pop();
    if (node->left_) q.push(node->left_);
    if (node->right_) q.push(node->right_);
    res.push_back(node->val_);
  }
}

std::vector<std::vector<int>> LevelOrderTraversal(BinaryTreeNode *root) {
  std::vector<std::vector<int>> res;
  if (root == nullptr) return res;
  std::queue<BinaryTreeNode *> q;
  q.push(root);
  while (!q.empty()) {
    auto qsize = q.size();
    std::vector<int> vec(qsize);
    for (int i = 0; i < qsize; i++) {  // level
      BinaryTreeNode *node = q.front();
      q.pop();
      if (node->left_) q.push(node->left_);
      if (node->right_) q.push(node->right_);
      vec[i] = node->val_;
    }
    res.push_back(vec);
  }
  return res;
}
